/*
 * memory/physical.c
 *
 * Copyright (C) 2018 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <memory.h>
#include <log.h>
#include <exception.h>
#include <cpu.h>

typedef struct
{
    page_t *all;
    page_num_t first_page_num;
    page_num_t num_pages;
    mini_list_entry_t free_stack;
    lock_t lock;
} page_list_t;

static page_list_t page_list[MAX_PHYS_ADDR_BITS - MIN_PHYS_ADDR_BITS + 1];

static inline byte_t count_bits(physical_t address)
{
    byte_t number = 64 - __builtin_clzll(address);
    ASSERT(number <= cpu_max_physical_bits);
    return number >= MIN_PHYS_ADDR_BITS ? number : MIN_PHYS_ADDR_BITS;
}

page_t *memory_acquire_page(byte_t min_bits, byte_t max_bits, size_t alignment)
{
    page_t *result = NULL;
    ASSERT(min_bits <= max_bits);
    ASSERT(alignment != 0);

    if (min_bits < MIN_PHYS_ADDR_BITS) min_bits = MIN_PHYS_ADDR_BITS;
    if (max_bits > cpu_max_physical_bits) max_bits = cpu_max_physical_bits;

    byte_t bits;
    for (bits = max_bits; bits >= min_bits && result == NULL; bits--)
    {
        page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];
        lock_acquire_smart(&list->lock);

        mini_list_entry_t *ptr;
        for (ptr = list->free_stack.next; ptr != &list->free_stack; ptr = ptr->next)
        {
            page_t *page = CONTAINER_OF(ptr, page_t, stack_link);
            if (page->number % PAGE_NUMBER(alignment) == 0)
            {
                mini_list_remove(&page->stack_link);
                page->status = PAGE_STATUS_ALLOCATED;
                result = page;
                break;
            }
        }

        lock_release(&list->lock);
    }

    return result;
}

void memory_acquire_area(byte_t min_bits, byte_t max_bits, size_t size, size_t alignment, area_t *area)
{
    ASSERT(size != 0);
    ASSERT(alignment != 0);
    ASSERT(min_bits <= max_bits);

    if (min_bits < MIN_PHYS_ADDR_BITS) min_bits = MIN_PHYS_ADDR_BITS;
    if (max_bits > cpu_max_physical_bits) max_bits = cpu_max_physical_bits;

    if (size == 1)
    {
        area->pages = memory_acquire_page(min_bits, max_bits, alignment);
        area->count = 1;
        return;
    }

    size_t num_pages = PAGE_NUMBER(PAGE_ALIGN_UP(size));
    area->pages = NULL;
    area->count = 0;

    byte_t bits;
    for (bits = max_bits; bits >= min_bits && area->pages == NULL; bits--)
    {
        page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];
        lock_acquire_smart(&list->lock);

        page_t *last_page = &list->all[list->num_pages - 1];
        while (!area->pages && last_page >= list->all)
        {
            for (; last_page >= list->all; last_page--)
            {
                if (memory_get_page_mapping(memory_default_table, last_page))
                {
                    if (last_page->status == PAGE_STATUS_FREE
                        && ((last_page->number - size) % PAGE_NUMBER(alignment) == 0)) break;
                }
                else
                {
                    last_page = ((page_t*)PAGE_ALIGN((uintptr_t)last_page));
                }
            }

            if (last_page < list->all) break;

            page_t *first_page;
            for (first_page = last_page; first_page >= list->all; first_page--)
            {
                if (!memory_get_page_mapping(memory_default_table, first_page)
                    || first_page->status != PAGE_STATUS_FREE) break;

                if ((last_page - first_page) == num_pages)
                {
                    area->pages = first_page;
                    area->count = num_pages;
                    break;
                }
            }
        }

        if (area->pages != NULL)
        {
            page_num_t i;
            for (i = 0; i < area->count; i++)
            {
                area->pages[i].status = PAGE_STATUS_ALLOCATED;
                mini_list_remove(&area->pages[i].stack_link);
            }
        }

        lock_release(&list->lock);
    }
}

void memory_release_page(page_t *page)
{
    byte_t bits = count_bits(page->number * PAGE_SIZE);
    ASSERT(bits >= MIN_PHYS_ADDR_BITS && bits <= cpu_max_physical_bits);

    page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];
    lock_acquire_smart(&list->lock);

    ASSERT(page->status >= PAGE_STATUS_ALLOCATED);
    mini_list_prepend(&list->free_stack, &page->stack_link);
    page->status = PAGE_STATUS_FREE;

    lock_release(&list->lock);
}

void memory_release_area(const area_t *area)
{
    size_t i;
    for (i = 0; i < area->count; i++) memory_release_page(&area->pages[i]);
}

page_t *memory_find_page_by_address(physical_t address)
{
    byte_t bits = count_bits(address);
    ASSERT(bits >= MIN_PHYS_ADDR_BITS && bits <= cpu_max_physical_bits);

    page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];
    page_t *page = &list->all[PAGE_NUMBER(address) - list->first_page_num];

    memory_flags_t access_flags;
    return (memory_query_page_flags(memory_default_table, page, &access_flags) == ERR_SUCCESS
            && page->status != PAGE_STATUS_ABANDONED) ? page : NULL;
}

void memory_claim_physical_region(physical_t address, qword_t size, page_status_t initial_status)
{
    const page_num_t first_page = PAGE_NUMBER(address);
    const page_num_t last_page = PAGE_NUMBER(PAGE_ALIGN(address + size - 1));

    byte_t bits;
    byte_t min_bits = count_bits((physical_t)first_page * PAGE_SIZE);
    byte_t max_bits = count_bits((physical_t)last_page * PAGE_SIZE);

    for (bits = min_bits; bits <= max_bits; bits++)
    {
        page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];
        lock_acquire_smart(&list->lock);

        page_num_t local_first = list->first_page_num;
        page_num_t local_last = list->first_page_num + list->num_pages - 1;
        if (first_page > local_first) local_first = first_page;
        if (last_page < local_last) local_last = last_page;

        page_num_t number = local_first;
        while (number <= local_last)
        {
            page_t *page = &list->all[number - list->first_page_num];
            memory_flags_t access_flags;

            if (memory_query_page_flags(memory_default_table, page, &access_flags) != ERR_SUCCESS)
            {
                page_t temporary[MAX_PAGING_LEVELS];
                page_num_t num_temp = 0;

                for (;;)
                {
                    bool_t used = FALSE;
                    page_t *backing_page = memory_acquire_page(MIN_PHYS_ADDR_BITS, MAX_PHYS_ADDR_BITS, PAGE_SIZE);

                    if (!backing_page)
                    {
                        ASSERT(initial_status == PAGE_STATUS_FREE);
                        ASSERT(num_temp < MAX_PAGING_LEVELS);

                        backing_page = &temporary[num_temp];
                        backing_page->number = number + num_temp;
                        backing_page->status = PAGE_STATUS_ALLOCATED;
                        backing_page->map_count = 0;
                        num_temp++;
                        used = TRUE;
                    }

                    void *location = (void*)PAGE_ALIGN((uintptr_t)page);
                    sysret_t ret = memory_map_page(memory_default_table,
                                                   backing_page,
                                                   location,
                                                   MEMORY_FLAG_ACCESSIBLE | MEMORY_FLAG_WRITABLE | MEMORY_FLAG_STICKY);
                    if (ret == ERR_SUCCESS)
                    {
                        memset(location, 0, PAGE_SIZE);
                        break;
                    }

                    ASSERT(ret == ERR_NOMEMORY);

                    if (!used)
                    {
                        memory_release_page(backing_page);
                        ASSERT(initial_status == PAGE_STATUS_FREE);
                        ASSERT(num_temp < MAX_PAGING_LEVELS);

                        backing_page = &temporary[num_temp];
                        backing_page->number = number + num_temp;
                        backing_page->map_count = 0;
                        num_temp++;
                    }

                    backing_page->status = PAGE_STATUS_FREE;
                    mini_list_prepend(&list->free_stack, &backing_page->stack_link);
                }

                page_num_t i;
                for (i = 0; i < num_temp; i++)
                {
                    ASSERT(temporary[i].status != PAGE_STATUS_FREE);
                    page[i] = temporary[i];
                }

                number += num_temp;
            }

            if (page->status == PAGE_STATUS_ABANDONED)
            {
                page->number = number;
                page->status = initial_status;
                page->map_count = 0;
                if (initial_status == PAGE_STATUS_FREE) mini_list_prepend(&list->free_stack, &page->stack_link);
            }

            number++;
        }

        lock_release(&list->lock);
    }
}

void memory_abandon_physical_region(physical_t address, qword_t size)
{
    // TODO: Unimplemented
    // THINK: Is it really necessary? It wouldn't free up much memory...
}

static void memory_init_page_lists(void)
{
    ASSERT(cpu_max_physical_bits <= MAX_PHYS_ADDR_BITS);

    for (byte_t bits = cpu_max_physical_bits; bits >= MIN_PHYS_ADDR_BITS; bits--)
    {
        page_list_t *list = &page_list[bits - MIN_PHYS_ADDR_BITS];

        list->first_page_num = bits > MIN_PHYS_ADDR_BITS ? PAGE_NUMBER(1ULL << (bits - 1)) : 0;
        list->num_pages = PAGE_NUMBER(1ULL << bits) - list->first_page_num;
        list->all = memory_request_metadata_space(list->num_pages, sizeof(page_t));
        mini_list_init(&list->free_stack);
        lock_init(&list->lock);
    }
}

void memory_init_physical(memory_map_entry_t *mmap, size_t entry_count)
{
    const char *status_text[] = { "Not Usable", "Usable", "In Use", "Reserved" };
    memory_init_page_lists();

    log_write(LOG_NORMAL, "Memory map:\nBase\t\t\tLength\t\t\tType\n");
    log_write(LOG_NORMAL, "------------------------------------------------------------\n");

    size_t i;
    for (i = 0; i < entry_count; i++)
    {
        log_write(LOG_NORMAL, "0x%016llX\t0x%016llX\t%s\n", mmap[i].address, mmap[i].length, status_text[mmap[i].status]);
        if (mmap[i].status == PAGE_STATUS_FREE) memory_claim_physical_region(mmap[i].address, mmap[i].length, mmap[i].status);
    }

    log_write(LOG_NORMAL, "------------------------------------------------------------\n");

    for (i = 0; i < entry_count; i++)
    {
        if (mmap[i].status != PAGE_STATUS_FREE) memory_claim_physical_region(mmap[i].address, mmap[i].length, mmap[i].status);
    }
}
